import itertools
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.preprocessing import KBinsDiscretizer


# 函数功能：根据给定数据集和项集元素个数生成候选项集
# 参数说明：
# - data：事务数据集，是一个包含多个事务（每个事务为元素列表）的列表
# - k：表示要生成的候选项集中元素的个数
# 返回值：由符合要求的项集组成的列表
def create_candidates(data, k):
    candidate_set = set()
    for transaction in data:
        # 对每个事务内的元素进行排序后再生成组合，确保组合的有序性
        sorted_transaction = sorted(transaction)
        combinations = itertools.combinations(sorted_transaction, k)
        for combination in combinations:
            candidate_set.add(combination)
    return list(candidate_set)


# 函数功能：计算给定候选项集在数据集中的支持度
# 参数说明：
# - data：事务数据集，是一个包含多个事务（每个事务为元素列表）的列表
# - candidates：要计算支持度的候选项集列表，每个候选项集为元组形式
# 返回值：一个字典，键为候选项集，值为对应的支持度（浮点数）
def compute_support(data, candidates):
    total_num_transactions = len(data)
    support_dict = {}
    for candidate in candidates:
        count = 0
        for transaction in data:
            # 判断候选项集是否是事务的子集来统计出现次数
            if set(candidate).issubset(set(transaction)):
                count += 1
        support_dict[candidate] = count / total_num_transactions
    return support_dict


# 函数功能：基于Apriori算法思想，从给定数据集和最小支持度要求中找出频繁项集
# 参数说明：
# - data：事务数据集，是一个包含多个事务（每个事务为元素列表）的列表
# - min_support：最小支持度阈值，浮点数，只有支持度大于等于此值的项集才被视为频繁项集
# 返回值：包含各轮（不同元素个数对应的）频繁项集的列表，每个元素是一个字典，键为频繁项集，值为对应的支持度
def apriori(data, min_support):
    current_k = 1
    all_frequent_itemsets = []
    # 第一轮，生成单个元素的候选项集并计算支持度，筛选出频繁项集
    candidates = create_candidates(data, current_k)
    support_result = compute_support(data, candidates)
    frequent_itemsets = {item: sup for item, sup in support_result.items() if sup >= min_support}
    all_frequent_itemsets.append(frequent_itemsets)

    # 不断迭代，增加项集元素个数，生成新的候选项集并筛选频繁项集，直到没有新的频繁项集产生
    while frequent_itemsets:
        current_k += 1
        candidates = create_candidates([list(item) for item in frequent_itemsets], current_k)
        support_result = compute_support(data, candidates)
        frequent_itemsets = {item: sup for item, sup in support_result.items() if sup >= min_support}
        if frequent_itemsets:
            all_frequent_itemsets.append(frequent_itemsets)

    return all_frequent_itemsets


if __name__ == "__main__":
    # 加载鸢尾花数据集
    iris = load_iris()
    feature_df = pd.DataFrame(data=iris.data, columns=iris.feature_names)
    feature_df['species'] = iris.target

    # 对数值型特征进行离散化处理
    discretizer = KBinsDiscretizer(n_bins=3, encode='ordinal', strategy='uniform')
    discretized_df = pd.DataFrame(discretizer.fit_transform(feature_df[feature_df.columns[:-1]]))

    # 将离散化后的特征值转换为便于处理的字符串形式
    for col in discretized_df.columns:
        discretized_df[col] = discretized_df[col].apply(lambda x: f'{feature_df.columns[col]}_{int(x)}')

    # 将数据转换为事务列表格式
    transactions = discretized_df.apply(lambda x: x.tolist(), axis=1).tolist()

    min_support_value = 0.2  # 设置最小支持度阈值
    frequent_itemsets_result = apriori(transactions, min_support_value)

    # 输出各轮找到的频繁项集及其支持度
    for k, itemsets in enumerate(frequent_itemsets_result, start=1):
        print(f"频繁{k}-项集:")
        for itemset, support in itemsets.items():
            print(f"  {itemset}: 支持度 = {support:.2f}")